#include "undistortion.hpp"
#include <opencv2/opencv.hpp>
#include <cmath>
#include <iostream>

float undistort_radius(float rd, float k1, float k2, int max_iter) {
    if (rd == 0.0f) return 0.0f;

    float ru = rd;
    for (int i = 0; i < max_iter; ++i) {
        float ru2 = ru * ru;
        float ru4 = ru2 * ru2;

        float f = ru * (1 + k1 * ru2 + k2 * ru4) - rd;
        float df = 1 + 3 * k1 * ru2 + 5 * k2 * ru4;

        float delta = f / df;
        ru -= delta;

        if (std::abs(delta) < 1e-6f) break;
    }

    return ru;
}

void undistort_image(const cv::Mat& img, cv::Mat& result, float k1, float k2) {
    int w = img.cols, h = img.rows;
    cv::Point2f center(w / 2.0f, h / 2.0f);
    cv::Mat map_x(h, w, CV_32FC1), map_y(h, w, CV_32FC1);

    for (int y = 0; y < h; ++y) {
        for (int x = 0; x < w; ++x) {
            float nx = (x - center.x) / center.x;
            float ny = (y - center.y) / center.y;
            float rd = std::sqrt(nx * nx + ny * ny);

            float ru = undistort_radius(rd, k1, k2);
            float scale = (rd == 0.0f) ? 1.0f : ru / rd;

            float src_x = scale * nx * center.x + center.x;
            float src_y = scale * ny * center.y + center.y;

            map_x.at<float>(y, x) = src_x;
            map_y.at<float>(y, x) = src_y;
        }
    }

    cv::remap(img, result, map_x, map_y, cv::INTER_LINEAR);
}

cv::Point2f undistort_point(const cv::Point2f& p, float k1, float k2) {
    float r2 = p.x * p.x + p.y * p.y;
    float scale = 1.0f / (1.0f + k1 * r2 + k2 * r2 * r2);
    return p * scale;
}

float compute_total_error(const std::vector<cv::Point2f>& points, float k1, float k2) {
    std::vector<cv::Point2f> corrected;
    for (const auto& p : points) corrected.push_back(undistort_point(p, k1, k2));

    cv::Vec4f line;
    cv::fitLine(corrected, line, cv::DIST_L2, 0, 1e-2, 1e-2);
    float vx = line[0], vy = line[1], x0 = line[2], y0 = line[3];

    float a = -vy, b = vx, c = -(a * x0 + b * y0);
    float error = 0.0f;
    for (const auto& p : corrected) {
        float d = std::abs(a * p.x + b * p.y + c) / std::sqrt(a * a + b * b);
        error += d * d;
    }
    return error / corrected.size();
}

std::pair<float, float> optimize_k1_k2(const std::vector<cv::Point2f>& points) {
    float k1 = 0.0f, k2 = 0.0f;
    float lr = 1e-3;
    const int max_iter = 200;
    float prev_error = compute_total_error(points, k1, k2);

    for (int iter = 0; iter < max_iter; ++iter) {
        float h = 1e-5;
        float e = compute_total_error(points, k1, k2);
        float e1 = compute_total_error(points, k1 + h, k2);
        float e2 = compute_total_error(points, k1, k2 + h);

        float g1 = (e1 - e) / h;
        float g2 = (e2 - e) / h;

        if (e > prev_error) lr *= 0.5f;
        else lr *= 1.05f;

        prev_error = e;

        k1 -= lr * g1;
        k2 -= lr * g2;

        if (std::abs(g1) < 1e-8 && std::abs(g2) < 1e-8) break;

        if (iter % 10 == 0) {
            std::cout << "Iter: " << iter << ", Error: " << e << ", k1: " << k1 << ", k2: " << k2 << std::endl;
        }
    }

    return {k1, k2};
}
